from typing import Annotated
from langchain_core.tools import tool, InjectedToolCallId
from langgraph.prebuilt import create_react_agent, InjectedState
from langgraph.graph import StateGraph, START, MessagesState
from langgraph.types import Command
from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage, AIMessage

from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

@tool("transfer_to_hotel_assistant", description="Transfer to hotel assistant.")
def transfer_to_hotel_assistant(state: Annotated[MessagesState, InjectedState]):
    """Transfer to hotel assistant."""
    return Command(
        # name of the agent (node) to go to
        goto="hotel_assistant",
        # data to send to the agent
        update={"messages": [AIMessage(content="Successfully transferred to hotel assistant")]},
        # indicate to LangGraph that we need to navigate to
        # agent node in a parent graph
        graph=Command.PARENT,
    )

@tool("transfer_to_flight_assistant", description="Transfer to flight assistant.")
def transfer_to_flight_assistant(state: MessagesState):
    """Transfer to flight assistant."""
    return Command(
        # name of the agent (node) to go to
        goto="flight_assistant",
        # data to send to the agent
        update={"messages": [AIMessage(content="Successfully transferred to flight assistant")]},
        # indicate to LangGraph that we need to navigate to
        # agent node in a parent graph
        graph=Command.PARENT,
    )

# Simple agent tools
def book_hotel(hotel_name: str):
    """Book a hotel"""
    return f"Successfully booked a stay at {hotel_name}."

def book_flight(from_airport: str, to_airport: str):
    """Book a flight"""
    return f"Successfully booked a flight from {from_airport} to {to_airport}."

# Define agents
flight_assistant = create_react_agent(
    model=llm,
    tools=[book_flight, transfer_to_hotel_assistant],
    # tools=[book_flight],
    prompt="You are a flight booking assistant, book flights for the user first. If the user needs help with anything else, transfer to the hotel assistant. Do not book hotels. Do not ask for confirmation. Do not ask for additional information. Do not say anything else. Do not repeat the user's message. Do not say anything other than the booking confirmation. Do not say anything other than the transfer confirmation. Do not say anything other than the booking confirmation and the transfer confirmation. Do not say anything other than the booking confirmation, the transfer confirmation",
    name="flight_assistant"
)
hotel_assistant = create_react_agent(
    model=llm,
    tools=[book_hotel, transfer_to_flight_assistant],
    # tools=[book_hotel],
    prompt="You are a hotel booking assistant",
    name="hotel_assistant"
)

# Define multi-agent graph
multi_agent_graph = (
    StateGraph(MessagesState)
    .add_node(flight_assistant)
    .add_node(hotel_assistant)
    .add_edge(START, "flight_assistant")
    .compile()
)

# Run the multi-agent graph
for chunk in multi_agent_graph.stream(
    {
        "messages": [
            {
                "role": "user",
                "content": "Please book a flight from BOS to JFK and a stay at McKittrick Hotel"
            }
        ]
    }
):
    print("------------------------------------------------------------------------------------")
    if "flight_assistant" in chunk and chunk["flight_assistant"] and chunk["flight_assistant"]["messages"]:
        print("--------------------------------  flight_assistant   -------------------------------")
        messages = chunk["flight_assistant"]["messages"]
        for msg in messages:
            print(f">>>>>>>>>>> {msg.__class__.__name__}")
            print(f"{msg}")
            print("\n")
    elif "transfer_to_hotel_assistant" in chunk and chunk["transfer_to_hotel_assistant"] and chunk["transfer_to_hotel_assistant"]["messages"]:
        print("--------------------------------  transfer_to_hotel_assistant   -------------------------------")
        messages = chunk["transfer_to_hotel_assistant"]["messages"]
        for msg in messages:
            print(f">>>>>>>>>>> {msg.__class__.__name__}")
            print(f"{msg}")
            print("\n")
    elif "hotel_assistant" in chunk and chunk["hotel_assistant"] and chunk["hotel_assistant"]["messages"]:
        print("--------------------------------  hotel_assistant   -------------------------------")
        messages = chunk["hotel_assistant"]["messages"]
        for msg in messages:
            print(f">>>>>>>>>>> {msg.__class__.__name__}")
            print(f"{msg}")
            print("\n")
    else:
        print(chunk)